package cn.micro.ttl.mdc;

import cn.micro.core.context.CurrentContext;
import cn.micro.core.util.StringUtil;
import cn.micro.core.util.TraceIdUtil;
import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.atomic.AtomicInteger;

import static cn.micro.core.constant.Cants.CURRENT_TRACE_ID;
import static cn.micro.core.constant.Cants.X_CURRENT_REQUEST_ID;

public class MDCThreadWrap {

    public static Runnable runnable(Runnable runnable) {
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        return () -> {
            try {
                setContextMap(contextMap);
                runnable.run();
            } finally {
                clear();
            }
        };
    }

    public static <T> Callable<T> callable(Callable<T> callable) {
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        return () -> {
            try {
                setContextMap(contextMap);
                return callable.call();
            } finally {
                clear();
            }
        };
    }

    private static void setContextMap(Map<String, String> contextMap) {
        if (contextMap != null) {
            MDC.setContextMap(contextMap);
        }
        if (!CurrentContext.existsKey(CURRENT_TRACE_ID)) {
            String traceId = TraceIdUtil.genTraceId();
            MDC.put(X_CURRENT_REQUEST_ID, "[" + traceId + "]");
            CurrentContext.set(CURRENT_TRACE_ID, traceId);
            CurrentContext.set(traceId, new AtomicInteger(0));
        }
    }

    private static void clear() {
        MDC.clear();
        String traceId = CurrentContext.get(CURRENT_TRACE_ID);
        if (StringUtil.isNotBlank(traceId)) {
            CurrentContext.removeKey(traceId);
            CurrentContext.removeKey(CURRENT_TRACE_ID);
        }
    }

}
